from typing import Dict, Any, Optional
from src.core.output_manager_base import OutputManagerBase, SimulationParameters, MPI, fem, FENICSX_AVAILABLE
import os

class VTKOutputHandler(OutputManagerBase):
    """
    用于处理 VTX/VTK 文件输出的管理器。
    """
    def __init__(self, 
                 comm: Optional[MPI.Comm], 
                 params: SimulationParameters, 
                 fem_entities: Dict[str, Any], 
                 output_dir: str,
                 filename_prefix: str = "solution",
                 output_frequency: int = 1):
        """
        初始化 VTKOutputHandler。

        参数:
            comm (Optional[MPI.Comm]): MPI 通信器。
            params (SimulationParameters): 模拟参数。
            fem_entities (Dict[str, Any]): FEM 实体，如解函数。
            output_dir (str): 输出目录。
            filename_prefix (str): VTX 文件名的前缀。
            output_frequency (int): 输出频率 (每 N 次迭代/步骤输出一次)。
        """
        super().__init__(comm, params, fem_entities, output_dir)
        self.filename_prefix = filename_prefix
        self.output_frequency = output_frequency
        # 假设要输出的解函数在 fem_entities 中以 "solution_functions" 为键存储，值为列表
        # 或者，可以更具体地指定要输出的函数名称
        self.functions_to_output: list[fem.Function] = [] # type: ignore
        if FENICSX_AVAILABLE:
            # 示例：尝试从 fem_entities 获取所有 fem.Function 对象
            # 在实际应用中，可能需要更复杂的逻辑来选择要输出的函数
            for key, entity in self.fem_entities.items():
                if isinstance(entity, fem.Function): # type: ignore
                    self.functions_to_output.append(entity)
                    if self.comm is None or self.comm.rank == 0: # type: ignore
                        print(f"信息: VTKOutputHandler 将尝试输出函数: '{entity.name if hasattr(entity, 'name') else key}'") # type: ignore
            if not self.functions_to_output and (self.comm is None or self.comm.rank == 0): # type: ignore
                 print(f"警告: VTKOutputHandler 在 fem_entities 中未找到可输出的 fem.Function 对象。")
        
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print(f"VTKOutputHandler: 将以频率 {self.output_frequency} 保存VTK文件，前缀为 '{self.filename_prefix}'。")

    def is_output_step(self, iteration_count: int, current_time: float) -> bool:
        """
        判断是否为输出步骤。
        """
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return False
        if self.output_frequency == 0:
            return iteration_count == 0 # 仅在迭代0输出（如果频率为0）
        return iteration_count % self.output_frequency == 0

    def execute_step(self, iteration_count: int, current_time: float) -> None:
        """
        执行 VTK 文件输出。
        """
        if not FENICSX_AVAILABLE:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print("VTKOutputHandler: FEniCSx 不可用，跳过 VTK 输出。")
            return
        
        if not self.functions_to_output:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print("VTKOutputHandler: 没有要输出的函数，跳过 VTK 输出。")
            return

        filepath = os.path.join(self.output_dir, f"{self.filename_prefix}_{iteration_count}.bp")
        os.makedirs(self.output_dir, exist_ok=True)
        
        try:
            from dolfinx.io import VTXWriter # type: ignore
            # VTXWriter 通常是并行的
            with VTXWriter(self.comm, filepath, self.functions_to_output, engine="BP4") as vtx: # type: ignore
                vtx.write(current_time)
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print(f"VTKOutputHandler: 在迭代 {iteration_count}, 时间 {current_time:.2f} 输出。文件已保存到: {filepath}")
        except ImportError:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print("警告: VTKOutputHandler 无法导入 dolfinx.io.VTXWriter，无法保存 VTK 文件。")
        except Exception as e:
            if self.comm is None or self.comm.rank == 0: # type: ignore
                print(f"错误: VTKOutputHandler 保存 VTK 文件时发生错误: {e}")

    def finalize(self) -> None:
        """
        完成 VTK 输出操作。
        """
        if self.comm is None or self.comm.rank == 0: # type: ignore
            print("VTKOutputHandler: 完成所有输出操作。") 